from stockInfo import StockInfo
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dropout, Dense, LSTM
import matplotlib.pyplot as plt
import os
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error
import math


class TrainModel:
    stock_list = []
    stock2Scaler = dict()
    train_date = []
    valid_date = []
    num_to_input = 60
    num_to_pred = 10

    def __init__(self, stock_list, train_date, valid_date, num_to_input=60, num_to_pred=10):
        self.stock_list = stock_list
        self.stock2Scaler = dict.fromkeys(self.stock_list)
        self.train_date = train_date
        self.valid_date = valid_date
        self.num_to_input = num_to_input
        self.num_to_pred = num_to_pred

    def train(self):
        x_train = []
        y_train = []
        x_valid = []
        y_valid = []
        for stock_id in self.stock_list:
            si = StockInfo(stock_id)
            single_x_train = si.get_all_info_of_one_stack(start_date=self.train_date[0],
                                                          end_date=self.train_date[1]).values
            single_x_valid = si.get_all_info_of_one_stack(start_date=self.valid_date[0],
                                                          end_date=self.valid_date[1]).values

            # 归一化
            sc = MinMaxScaler(feature_range=(0, 1))  # 定义归一化：归一化到(0，1)之间
            single_x_train = sc.fit_transform(single_x_train)  # 求得训练集的最大值，最小值这些训练集固有的属性，并在训练集上进行归一化
            single_x_valid = sc.transform(single_x_valid)  # 利用训练集的属性对测试集进行归一化
            self.stock2Scaler[stock_id] = sc

            single_y_train = single_x_train[:, 1]
            single_y_valid = single_x_valid[:, 1]

            num_train_days = single_x_train.shape[0]
            num_valid_days = single_x_valid.shape[0]
            for i in range(self.num_to_input, num_train_days - self.num_to_pred):
                tmp_x_train = single_x_train[i - self.num_to_input: i, :]
                tmp_y_train = single_y_train[i: i + self.num_to_pred]
                # 空值处理
                if sum(sum(np.isnan(tmp_x_train))) > 0 or sum(np.isnan(tmp_y_train)) > 0:
                    continue
                x_train.append(tmp_x_train)
                y_train.append(tmp_y_train)

            for i in range(self.num_to_input, num_valid_days - self.num_to_pred):
                tmp_x_valid = single_x_valid[i - self.num_to_input: i, :]
                tmp_y_valid = single_y_valid[i: i + self.num_to_pred]
                # 空值处理
                if sum(sum(np.isnan(tmp_x_valid))) > 0 or sum(np.isnan(tmp_y_valid)) > 0:
                    continue
                x_valid.append(tmp_x_valid)
                y_valid.append(tmp_y_valid)

        np.random.seed(7)
        np.random.shuffle(x_train)
        np.random.seed(7)
        np.random.shuffle(y_train)

        x_train, y_train = np.array(x_train).astype(np.float32), np.array(y_train).astype(np.float32)
        x_valid, y_valid = np.array(x_valid).astype(np.float32), np.array(y_valid).astype(np.float32)

        # 使x_train符合RNN输入要求：[送入样本数， 循环核时间展开步数， 每个时间步输入特征个数]。
        # 此处整个数据集送入，送入样本数为x_train.shape[0]组数据；输入60个训练数据，预测出未来7天的收盘价，
        # 循环核时间展开步数为60; 每个时间步送入的特征是某一天的因子数据，故每个时间步输入特征个数为1
        # x_train = np.reshape(x_train, (x_train.shape[0], num_to_input, x_train.shape[2]))

        model = tf.keras.Sequential([
            LSTM(128, return_sequences=True),
            Dropout(0.5),
            LSTM(128),
            Dropout(0.5),
            # Dense(64, kernel_initializer="uniform", activation='relu'),
            Dense(self.num_to_pred, kernel_initializer='uniform', activation='relu')
        ])

        model.compile(optimizer=tf.keras.optimizers.Adam(3e-4),
                      loss='mean_squared_error')  # 损失函数用均方误差
        # 该应用只观测loss数值，不观测准确率，所以删去metrics选项，一会在每个epoch迭代显示时只显示loss值

        checkpoint_save_path = "checkpoint\\LSTM_stock.ckpt"

        cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                         save_weights_only=False,
                                                         save_best_only=True,
                                                         monitor='val_loss')

        history = model.fit(x_train, y_train, batch_size=128, epochs=20, validation_data=(x_valid, y_valid),
                            validation_freq=1,
                            callbacks=[cp_callback])

        model.summary()

        loss = history.history['loss']
        val_loss = history.history['val_loss']

        plt.plot(loss, label='Training Loss')
        plt.plot(val_loss, label='Validation Loss')
        plt.title('Training and Validation Loss')
        plt.legend()
        plt.show()

        return self.stock2Scaler


if __name__ == '__main__':

    stock_list = ['600150.XSHG', '600900.XSHG', '600048.XSHG', '600340.XSHG',
                  '600569.XSHG', '600019.XSHG', '600115.XSHG', '600118.XSHG',
                  '600151.XSHG', '002558.XSHE', '603110.XSHG', '000001.XSHE',
                  '600030.XSHG', '601001.XSHG', '002611.XSHE', '601857.XSHG',
                  '600028.XSHG', '600050.XSHG', '000063.XSHE', '000651.XSHE',
                  '600111.XSHG', '600518.XSHG', '600056.XSHG', '600519.XSHG']

    train_date = ['2010-01-04', '2015-12-31']
    valid_date = ['2016-01-01', '2016-12-31']
    num_to_input = 60
    num_to_pred = 10

    tm = TrainModel(stock_list, train_date, valid_date)
    stock2Scaler = tm.train()   # 执行训练过程，并且后续要用到反归一化


